#ifndef FITNESS_HPP_INCLUDED
#define FITNESS_HPP_INCLUDED

#include <vector>
#include <math.h>

#define NOBACKGROUND -999.0

//
// Calculate the fitness given the genotype and selection coefficient
float getDipFitness(std::vector<bool> &hap1, std::vector<bool> &hap2, std::vector<float> &effects)
{
    float thisFitness = 1.0;
    for(unsigned int i=0; i<hap1.size(); i++) // each marker
    {
        if(effects[i] > 0.0) // reference state is favored
        {
            // both are at reference state
            if(!hap1[i] and !hap2[i]) thisFitness = thisFitness * pow(1.0 + effects[i],2.0);
            // heterozygote
            else if(hap1[i] xor hap2[i]) thisFitness = thisFitness * (1.0 + effects[i]);
        }
        else if(effects[i] < 0.0) // variant state is favored
        {
            // both are at variant state
            if(hap1[i] and hap2[i]) thisFitness = thisFitness * pow(1.0 - effects[i],2);
            // heterozygote
            else if(hap1[i] xor hap2[i]) thisFitness = thisFitness * (1.0 - effects[i]);
        }
    }
    return thisFitness;
}

float getHapFitness(std::vector<bool> &hap, std::vector<float> &effects)
{
    float thisFitness = 1.0;
    for(unsigned int i=0; i<hap.size(); i++) // each marker
    {
        if(effects[i] > 0.0) // reference state is favored
        {
            // at reference state
            if(!hap[i]) thisFitness = thisFitness * (1.0 + effects[i]);
        }
        else if(effects[i] < 0.0) // variant state is favored
        {
            // at variant state
            if(hap[i]) thisFitness = thisFitness * (1.0 - effects[i]);
        }
    }
    return thisFitness;
}

float getHapIntFitness(std::vector<bool> &hap, std::vector<float> &effects)
{
    float thisFitness = 1.0;
    float intEffect = 0.0;
    for(unsigned int i=0; i<hap.size(); i++) // each marker
    {
        if( (effects[i] > 0.0) & (effects[i] < 100.0) )  // reference state is favored
        {
            // at reference state
            if(!hap[i]) thisFitness = thisFitness * (1.0 + effects[i]);
        }
        else if( (effects[i] < 0.0) & (effects[i] > -100.0) ) // variant state is favored
        {
            // at variant state
            if(hap[i]) thisFitness = thisFitness * (1.0 - effects[i]);
        }
        if( fabs(effects[i]) >= 100 ) // an interacting loci
            // a site with effect = 100 will have its reference allele as the background for another site
            // a site with effect = -100 will have its alternative allele as the background for another site
        {
            if(intEffect == 0.0)
            {
                if( effects[i] >= 100.0 )  // reference state is favored OR as background
                {
                    // at reference state
                    if (!hap[i])
                    {
                        intEffect = effects[i];
                    }
                    else
                    {
                        intEffect = NOBACKGROUND;
                    }
                }
                else if( effects[i] <= -100.0 ) // variant state is favored OR as background
                {
                    // at variant state
                    if(hap[i])
                    {
                        intEffect = effects[i];
                    }
                    else
                    {
                        intEffect = NOBACKGROUND;
                    }
                }
            }
            else if(intEffect - NOBACKGROUND > 1.0)
            {
                if( effects[i] >= 100.0 )  // reference state is favored OR as background
                {
                    // at reference state
                    if(!hap[i])
                    {
                        thisFitness = thisFitness * (1.0 + fabs(fabs(intEffect) - fabs(effects[i]) ));
                    }
                }
                else if( effects[i] <= -100.0 ) // variant state is favored OR as background
                {
                    // at variant state
                    if(hap[i])
                    {
                        thisFitness = thisFitness * (1.0 + fabs( fabs(intEffect) - fabs(effects[i])) );
                    }
                }
            }
        }
    }
    return thisFitness;
}


#endif // FITNESS_HPP_INCLUDED
